// STL
#include <fstream>
#include <iostream>
#include <numeric>
// BOOST
#include <boost/mpi.hpp>
#include <boost/program_options.hpp>
#include <boost/iostreams/filtering_streambuf.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/filter/bzip2.hpp>
// ARMADILLO
#include <armadillo>
// FACTORIZATION
#include "helper/helper.h"
#include "helper/mpi_helper.h"
#include "helper/randGen.h"
#include "factor_jobs.h"

using namespace std;
using namespace arma;
namespace mpi = boost::mpi;
namespace po = boost::program_options;

class TEST {
public:
  mpi::communicator& world;
  po::variables_map& vm;

  TEST(mpi::communicator& world_, po::variables_map& vm_) 
    : world(world_), vm(vm_) {}
  void run_test();
  bool test_rank();
  bool test_arma();
  bool test_mpi();
  bool test_saveload();
  bool test_io();
  bool test_rand();
};

int main(int argc, char* argv[]) {
  mpi::environment env(argc, argv);
  mpi::communicator world;
  po::options_description desc("Allowed options");
  po::variables_map vm;
  JOB job(world, vm);
  TEST test(world, vm);

  job.processParameter(argc, argv);

  posix::ptime start_time(CLOCK);

  if (vm.count("help"))
    MSG(desc);
  else {
    REQUIRE_PARAM("job");
    string job_type = vm["job"].as<string>();
    if (job_type == "runtest")
      test.run_test();
    else if (job_type == "multNMF")
      job.doMultNMF();
    else if (job_type == "parMultNMF")
      job.doParMultNMF();
    else if (job_type == "hals")
      job.doHALS();
    else if (job_type == "parHals")
      job.doParHALS();
    else if (job_type == "stochHals")
      job.doStochHALS();
    else if (job_type == "parStochHals")
      job.doParStochHALS();
    else
      FATAL("Unknown job type")
  }

  posix::time_duration d(CLOCK-start_time);
  if (world.rank() == 0) MSG("d=" << SECONDS(d) << " sec");

  return 0;
}

void TEST::run_test() {
  //MSG("TEST_RANK = " << (test_rank() ? "PASSED":"FAILED"));
  //MSG("TEST_ARMA = " << (test_arma() ? "PASSED":"FAILED"));
  //MSG("TEST_MPI " << world.rank() << " = " << (test_mpi() ? "PASSED":"FAILED"));
  //MSG("TEST_SAVELOAD " << world.rank() << " = " << (test_saveload() ? "PASSED":"FAILED"));
  //MSG("TEST_IO " << world.rank() << " = " << (test_io() ? "PASSED":"FAILED"));
  MSG("TEST_RAND " << world.rank() << " = " << (test_rand() ? "PASSED":"FAILED"));
}

bool TEST::test_rank() {
  MSG("Hello, world !!!" << world.rank() << " of " << world.size());
  return true;
}

bool TEST::test_arma() {
  mat A = randu<mat>(10,10);
  vec b = randu<vec>(10);
  vec x = solve(A,b);
  MSG("rank=" << world.rank() << " x=" << endl << x);

  mat B = zeros<mat>(10,10);
  MSG("B+1=\n" << (B+1));
  return true;
}

bool TEST::test_mpi() {
  mat a = world.rank()*ones<mat>(5,5);
  MSG("a[" << world.rank() << "] ready");
  vector<mat> all_a;
  gather(world, a, all_a, 0);

  if (world.rank() == 0 && vm["verbose"].as<int>()) {
    //FORE(i, 0, SIZE(all_a)) MSG("a[" << i << "]=\n" << all_a[i]);
    mat a_big;
    PART<mat>::cat(all_a, a_big);
    //MSG("a_big=\n" << a_big);

    int n_part = 3;
    vector<int> r;
    vector<mat> a_split;
    PART<mat>::split_idx(a_big.n_rows, n_part, r);
    PART<mat>::split(a_big, r, a_split);
    //FORE(i, 0, SIZE(a_split)) MSG("a_s[" << i << "]=\n" << a_split[i]);
  }

  mat b;
  PART<mat>::broadcast(world, all_a, b);
  MSG("b[" << world.rank() << "]=\n" << b);

  return true;
}

bool TEST::test_saveload() {
  mat a = world.rank()*ones<mat>(5,5), b;
  PART<mat>::save(world, "data/A", a);
  PART<mat>::load(world, "data/A", b);
  MSG("b[" << world.rank() << "]=\n" << b);
  return true;
}

namespace io = boost::iostreams;

bool TEST::test_io() {
  mat W = randu<mat>(1000,20), H = randu<mat>(1200,20);
  mat A = W*trans(H);
  
  {
    ofstream fout("test.gz", ios::out | ios::binary);
    io::filtering_streambuf<io::output> out_buf;
    //out_buf.push(io::gzip_compressor());
    out_buf.push(io::bzip2_compressor());
    out_buf.push(fout);
    ostream out(&out_buf);
    out.write((const char*)A.memptr(), A.n_elem*sizeof(double));
    out.flush();
  }

  ifstream fin("test.gz", ios::in | ios::binary);
  io::filtering_streambuf<io::input> in_buf;
  //in_buf.push(io::gzip_decompressor());
  in_buf.push(io::bzip2_decompressor());
  in_buf.push(fin);
  istream in(&in_buf);
  double *v = new double[A.n_elem*2];
  in.read((char*)v, A.n_elem*2*sizeof(double));
  int n = in.gcount();
  MSG(n << " " << (n / sizeof(double)));
  n /= sizeof(double);
  if (n != (int)A.n_elem) return false;
  mat B(v, A.n_rows, A.n_cols);
  MSG(accu(A == B));
  return accu(A==B)==A.n_elem;
}

bool TEST::test_rand() {
  mat Arows, Acols;
  RAND<mat> ra(world, vm);
  ra.randNMF(Arows, Acols);
  MSG("mr=" << Arows.n_rows << " nc=" << Acols.n_cols 
      << " m=" << Acols.n_rows << " n=" << Arows.n_cols);
  int mr = Arows.n_rows, nc = Acols.n_cols;
  int m = Acols.n_rows, n = Arows.n_cols;
  vector<int> r,c;
  vector<mat> Ar;
  mat A;
  gather(world, mr, r, 0);
  gather(world, nc, c, 0);
  gather(world, Arows, Ar, 0);
  if (MPI_IS_MASTER) {
    PART<mat>::cat(Ar, A);
    return accumulate(ALL(r), 0) == m &&
      accumulate(ALL(c), 0) == n &&
      accu(A == ra.W*trans(ra.H)) == m*n ;
  }
  return true;
}
